import os
import shutil
import sys
from PIL import Image
import math
import skimage
import cv2
import numpy as np
from skimage.io import imread
from skimage import exposure

def add_gauss_noise(image, var):
    image_noise = skimage.util.random_noise(image, mode='gaussian', var = var)
    return image_noise

def add_salt_pepper_noise(image, amount):
    image_noise = skimage.util.random_noise(image, mode='s&p', amount = amount)
    return image_noise

def add_speckle_noise(image, var):
    image_noise = skimage.util.random_noise(image, mode='speckle', var = var)
    return image_noise

def add_median_filter(image, kernal):
    image_noise = cv2.medianBlur(image, kernal)
    return image_noise

def add_mean_filter(image, kernal):
    image_noise = cv2.blur(image, (kernal, kernal))
    return image_noise

def add_gauss_filter(image, kernal):
    image_noise = cv2.GaussianBlur(image, (kernal, kernal), 0)
    return image_noise

def add_CHE_filter(image):
    ycrcb_img = cv2.cvtColor(image, cv2.COLOR_BGR2YCrCb)
    ycrcb_img[:, :, 0] = cv2.equalizeHist(ycrcb_img[:, :, 0])
    image_noise = cv2.cvtColor(ycrcb_img, cv2.COLOR_YCrCb2BGR)
    return image_noise

def add_gamma_correction(image, gamma):
    # image_noise = np.power(image, 1 / gamma)
    image_noise = exposure.adjust_gamma(image, gamma)
    return image_noise

def rotate(image, angle):
    rows, cols, channel = image.shape
    M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1.0)
    image_noise = cv2.warpAffine(image, M, (cols, rows))
    return image_noise

def translate(image, x, y):
    rows, cols, channel = image.shape
    M = np.float32([[1, 0, x], [0, 1, y]])
    image_noise = cv2.warpAffine(image, M, (cols, rows))
    return image_noise

def scaling(image, factor):
    rows, cols, channel = image.shape
    # x3.0
    image_noise = cv2.resize(image, (factor * cols, factor * rows))
    # return
    image_noise = cv2.resize(image_noise, (cols, rows))
    return image_noise

def Jpegcompression(image, q):
    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), q]
    result, encimg = cv2.imencode('.jpg', image, encode_param)
    image_noise = cv2.imdecode(encimg, 1)
    return image_noise

def CenterCrop(image, percentage):
    rows, cols, channel = image.shape
    ratio = math.sqrt(percentage)
    y_begin = int(rows * (1 - ratio) / 2)
    y_end = y_begin + int(rows * ratio)
    x_begin = int(cols * (1 - ratio) / 2)
    x_end = x_begin + int(cols * ratio)
    CenterBlock = image[y_begin:y_end, x_begin:x_end]
    CenterBlock_padding = cv2.copyMakeBorder(CenterBlock, y_begin, (rows - y_end), x_begin, (cols - x_end), cv2.BORDER_CONSTANT, value = (0, 0, 0))
    image_noise = image - CenterBlock_padding
    return image_noise

def EdgeCrop(image, percentage):
    rows, cols, channel = image.shape
    center_percentage = 1 - percentage
    ratio = math.sqrt(center_percentage)
    y_begin = int(rows * (1 - ratio) / 2)
    y_end = y_begin + int(rows * ratio)
    x_begin = int(cols * (1 - ratio) / 2)
    x_end = x_begin + int(cols * ratio)
    CenterBlock = image[y_begin:y_end, x_begin:x_end]
    CenterBlock_padding = cv2.copyMakeBorder(CenterBlock, y_begin, (rows - y_end), x_begin, (cols - x_end), cv2.BORDER_CONSTANT, value = (0, 0, 0))
    image_noise = CenterBlock_padding
    return image_noise

def opencv_rotate(img, angle):
    h, w = img.shape[:2]
    center = (w / 2, h / 2)
    scale = 1.0
    M = cv2.getRotationMatrix2D(center, angle, scale)
    # 2.2 新的宽高，radians(angle) 把角度转为弧度 sin(弧度)
    new_H = int(w * math.fabs(math.sin(math.radians(angle))) + h * math.fabs(math.cos(math.radians(angle))))
    new_W = int(h * math.fabs(math.sin(math.radians(angle))) + w * math.fabs(math.cos(math.radians(angle))))
    # 2.3 平移
    M[0, 2] += (new_W - w) / 2
    M[1, 2] += (new_H - h) / 2
    rotate = cv2.warpAffine(img, M, (new_W, new_H), borderValue=(0, 0, 0))
    return rotate

current_path = os.getcwd()
# imageslist = current_path + '/our_test_data_robust/ori/VOC2012'
imageslist = current_path + '/our_test_data_robust/ori/Holidays_500images'  # '/our_test_data_robust/COCO2014'
# images_attacks1_list = current_path + '/Dataset_training/VOC2007/attacks/gauss_noise0.001'
index = 1
# for images_file_path in sorted(os.listdir(imageslist), key=lambda x: int(str(x).split('.')[0])):
for images_file_path in sorted(os.listdir(imageslist), key=lambda x: int(str(x).split('.')[0])):  # COCO2014: sorted(os.listdir(imageslist), key=lambda x: int(str(x).split('.')[0].split('_')[2]))
    # print(images_file_path)
    image_file_path_total = os.path.join(imageslist, images_file_path)
    image_file = cv2.imread(image_file_path_total)
    # image_file = Image.open(image_file_path)
    # add noise
    image_file_noise1 = add_gauss_noise(image_file, 0.001)
    image_file_noise2 = add_salt_pepper_noise(image_file, 0.001)
    image_file_noise3 = add_speckle_noise(image_file, 0.01)
    image_file_noise4 = add_median_filter(image_file, 3)
    image_file_noise5 = add_mean_filter(image_file, 3)
    image_file_noise6 = add_gauss_filter(image_file, 3)
    image_file_noise7 = add_CHE_filter(image_file)
    image_file_noise8 = add_gamma_correction(image_file, 0.8)
    image_file_noise9 = rotate(image_file, 10)
    image_file_noise10 = translate(image_file, 36, 20)

    # add new Faster-rcnn
    image_file_noise11 = rotate(image_file, 30)
    image_file_noise12 = rotate(image_file, 50)
    image_file_noise13 = translate(image_file, 40, 25)
    image_file_noise14 = translate(image_file, 80, 50)
    image_file_noise15 = scaling(image_file, 3)
    image_file_noise16 = Jpegcompression(image_file, 90)
    image_file_noise17 = CenterCrop(image_file, 0.05)
    image_file_noise18 = CenterCrop(image_file, 0.1)
    image_file_noise19 = CenterCrop(image_file, 0.2)
    image_file_noise20 = EdgeCrop(image_file, 0.05)
    image_file_noise21 = EdgeCrop(image_file, 0.1)
    image_file_noise22 = EdgeCrop(image_file, 0.2)



    save_dir1 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Gauss noise0.001')  # 'ori/COCO2014', 'COCO2014_cv2_skimage/Gauss noise0.001'
    save_dir2 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Salt and pepper noise0.001')
    save_dir3 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Speckle noise0.01')
    save_dir4 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Median filter3')
    save_dir5 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Mean filter3')
    save_dir6 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Gauss filter3')
    save_dir7 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/CHE')
    save_dir8 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Gamma correction0.8')
    save_dir9 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Rotation10')
    save_dir10 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Translation36')

    # add new Faster-rcnn
    save_dir11 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Rotation30')
    save_dir12 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Rotation50')
    save_dir13 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Translation40')
    save_dir14 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Translation80')
    save_dir15 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Scaling3')
    save_dir16 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Jpeg90')
    save_dir17 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/CenterCrop0.05')
    save_dir18 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/CenterCrop0.1')
    save_dir19 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/CenterCrop0.2')
    save_dir20 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/EdgeCrop0.05')
    save_dir21 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/EdgeCrop0.1')
    save_dir22 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/EdgeCrop0.2')

    cv2.imwrite(save_dir1, 255 * image_file_noise1)
    cv2.imwrite(save_dir2, 255 * image_file_noise2)
    cv2.imwrite(save_dir3, 255 * image_file_noise3)
    cv2.imwrite(save_dir4, image_file_noise4)
    cv2.imwrite(save_dir5, image_file_noise5)
    cv2.imwrite(save_dir6, image_file_noise6)
    cv2.imwrite(save_dir7, image_file_noise7)
    cv2.imwrite(save_dir8, image_file_noise8)
    cv2.imwrite(save_dir9, image_file_noise9)
    cv2.imwrite(save_dir10, image_file_noise10)

    # add new Faster-rcnn
    cv2.imwrite(save_dir11, image_file_noise11)
    cv2.imwrite(save_dir12, image_file_noise12)
    cv2.imwrite(save_dir13, image_file_noise13)
    cv2.imwrite(save_dir14, image_file_noise14)
    cv2.imwrite(save_dir15, image_file_noise15)
    cv2.imwrite(save_dir16, image_file_noise16)
    cv2.imwrite(save_dir17, image_file_noise17)
    cv2.imwrite(save_dir18, image_file_noise18)
    cv2.imwrite(save_dir19, image_file_noise19)
    cv2.imwrite(save_dir20, image_file_noise20)
    cv2.imwrite(save_dir21, image_file_noise21)
    cv2.imwrite(save_dir22, image_file_noise22)

    # print('ori:', image_file.shape)
    # image_file_noise9 = opencv_rotate(image_file, 10)
    # print('attacked:',image_file_noise9.shape)
    # save_dir9 = image_file_path_total.replace('ori/COCO2014', 'COCO2014_cv2_skimage/Rotation10_new')
    # cv2.imwrite(save_dir9, image_file_noise9)

    image_file_noise23 = Jpegcompression(image_file, 10)
    save_dir23 = image_file_path_total.replace('ori/Holidays_500images_middle', 'Holidays_500images_middle_cv2_skimage/Jpeg10')
    cv2.imwrite(save_dir23, image_file_noise23)
print('done')
